package cc.shacocloud.luanniao.core.service.ai

import org.springframework.ai.chat.messages.Message
import org.springframework.ai.chat.messages.UserMessage
import org.springframework.ai.chat.model.ChatResponse
import org.springframework.ai.chat.prompt.Prompt
import org.springframework.ai.document.Document
import org.springframework.ai.embedding.EmbeddingOptionsBuilder
import org.springframework.ai.embedding.EmbeddingRequest
import org.springframework.ai.embedding.EmbeddingResponse
import reactor.core.publisher.Flux

/**
 * ai 模型顶层接口
 *
 * @author 思追(shaco)
 */
interface AiModel {

    // ------------ 聊天相关

    /**
     * 单消息对话
     */
    fun chatCall(message: String): String {
        val prompt = Prompt(UserMessage(message))
        val generation = chatCall(prompt).result
        return if ((generation != null)) generation.output.content else ""
    }

    /**
     * 多消息对话
     */
    fun chatCall(messages: List<Message>): String {
        val prompt = Prompt(messages)
        val generation = chatCall(prompt).result
        return if ((generation != null)) generation.output.content else ""
    }

    /**
     * 对话
     */
    fun chatCall(prompt: Prompt): ChatResponse {
        throw UnsupportedOperationException()
    }

    /**
     *
     */
    fun chatStream(message: String): Flux<String> {
        val prompt = Prompt(message)
        return chatStream(prompt).map { response: ChatResponse ->
            if ((response.result == null || response.result.output == null || response.result.output.content == null))
                ""
            else
                response.result.output.content
        }
    }

    fun chatStream(messages: List<Message>): Flux<String> {
        val prompt = Prompt(messages)
        return chatStream(prompt).map { response: ChatResponse ->
            if ((response.result == null || response.result.output == null || response.result.output.content == null))
                ""
            else
                response.result.output.content
        }
    }

    /**
     * 流消息对话
     */
    fun chatStream(prompt: Prompt): Flux<ChatResponse> {
        throw UnsupportedOperationException()
    }

    // ------------ 向量相关

    fun embedCall(request: EmbeddingRequest): EmbeddingResponse {
        throw UnsupportedOperationException()
    }

    /**
     * 将给定的文本嵌入到向量中
     */
    fun embed(text: String): FloatArray {
        require(text.isNotBlank()) { "文本不可用为空" }
        val response = embed(listOf(text))
        return response[0]
    }

    /**
     * 将给定文档的内容嵌入到向量中
     */
    fun embed(document: Document): FloatArray {
        throw UnsupportedOperationException()
    }

    /**
     * 将一批文本嵌入到向量中
     */
    fun embed(texts: List<String>): List<FloatArray> {
        return embedForResponse(texts)
            .results
            .stream()
            .map { it.output }
            .toList()
    }

    /**
     * 将一批文本嵌入到向量中，并返回 [EmbeddingResponse]
     */
    fun embedForResponse(texts: List<String>): EmbeddingResponse {
        require(texts.isNotEmpty()) { "长度不可用为空" }
        return embedCall(EmbeddingRequest(texts, EmbeddingOptionsBuilder.builder().build()))
    }

    // -----------

    /**
     * 文本转为 token
     */
    fun textToTokens(text: String): List<Int> {
        return GPT2Tokenizer.DEFAULT.encode(text)
    }
}